#include "Aimbot.h"

Aimbot::Aimbot() = default;

void
Aimbot::run(cv::Mat &img, Armor::Color enemy_color, std::vector<Armor> &armors, int max_armor_size)
{
    armors.clear();
    armors.reserve(2 * max_armor_size);
    cv::Mat processed_img;
    std::vector<Armor::Lightbar> lightbars;
    processImage(img, enemy_color, processed_img);            /* 图像预处理 */
    findLightbars(img, processed_img, enemy_color, lightbars); /* 寻找灯条 */
    matchArmors(img, lightbars, armors);                       /* 灯条两两匹配装甲板 */
    gradeArmors(armors, max_armor_size);
}

void
Aimbot::findLightbars(const cv::Mat &img, const cv::Mat &processed_image, Armor::Color mcu_color,
                      std::vector<Armor::Lightbar> &lightbars)
{
    lightbars.clear();
    lightbars.reserve(30);
    // 找出所有轮廓
    std::vector<std::vector<cv::Point>> contours;
    std::vector<cv::Vec4i> hierarchy;
    findContours(processed_image, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);// 寻找最外层轮廓
    for (auto &contour: contours)
    {
        if (contour.size() < 6) continue;
        Armor::Lightbar lightbar;
        lightbar.rect = cv::minAreaRect(contour);
        lightbar.elli_rect = cv::fitEllipse(contour);
        // 面积筛选
        if (lightbar.rect.size.area() < MAX_LIGHTBAR_AREA
            && lightbar.rect.size.area() > MIN_LIGHTBAR_AREA)
        {
            //矫正灯条minAreaRect，保证 height > width
            adjustLightBarRect(lightbar.rect);
            adjustLightBarAngle(lightbar.elli_rect.angle, lightbar.angle);

            if (lightbar.rect.size.height / lightbar.rect.size.width < 1.5) continue;
            //角度、颜色筛选灯条
            if (abs(lightbar.angle) < MAX_LIGHTBAR_ANGLE)
            {
                judgeLightbarColor(lightbar, img);
                if (lightbar.color == mcu_color)
                { lightbars.emplace_back(lightbar); }
            }
        }
    }
    /* 对候选灯条进行排序 */
    sort(lightbars.begin(), lightbars.end(), Armor::Lightbar::coordinateComparator);
}

void
Aimbot::matchArmors(const cv::Mat &img, const std::vector<Armor::Lightbar> &lightbars,
                    std::vector<Armor> &armors)
{
// 注意要检查是否是empty，因为begin()和end()是不检查是否为空的
    if (!checkLightbarSize(lightbars)) return;

    for (auto lbar = lightbars.begin(); lbar != lightbars.end(); ++lbar)
    {
        auto rbar = lbar;
        for (++rbar; rbar != lightbars.end(); ++rbar)
        {
            //------------------像素距离-----------------------------------
            float delta_width = fabs(rbar->rect.center.x - lbar->rect.center.x);
            float delta_height = fabs(rbar->rect.center.y - lbar->rect.center.y);
            // 装甲板x距离过近
            if (delta_width < lbar->rect.size.height) continue;
            // 装甲板x距离过远
            if (delta_width > lbar->rect.size.height * (MAX_ASPECT_BIG_RATIO + 1)) break;
            // 装甲板y距离过远
            if (delta_height > 150.0) continue;

            //---------------装甲板比例--------------------------------------
            double armor_width = sqrt(delta_width * delta_width + delta_height * delta_height);
            double armor_height = (lbar->rect.size.height + rbar->rect.size.height) / 2;
            double armor_ratio = armor_width / armor_height;
            // 默认小装甲板
            bool isBigArmor = false;
            if (armor_ratio > 3.4)
            { isBigArmor = true; }
            if (armor_ratio < MIN_ASPECT_SMALL_RATIO || armor_ratio > MAX_ASPECT_BIG_RATIO)
            {
                continue;
            }

            //--------------装甲板三维空间距离------------------
            double armor_distance = CAMERA_fy * 0.055 / armor_height;
            if (armor_distance > 15.0)// 注意,哨兵不一样
                continue;

            //--------------装甲板倾斜角度-----------------------
            double armor_angle = fabs(atan2(delta_height, delta_width) * 180 / M_PI);
            if (armor_angle > MAX_ARMOR_ANGLE) continue;

            //--------------两灯条长度比-----------------------------------------
            double lbar_len_ratio = lbar->rect.size.height / rbar->rect.size.height;
            lbar_len_ratio = ((lbar_len_ratio > 1.0) ? lbar_len_ratio : (1.0 / lbar_len_ratio));
            if (lbar_len_ratio > MAX_LENGTH_RATIO) continue;

            //-------------两灯条角度差------------------------------------------------
            double lbar_delta_angle = abs(lbar->angle - rbar->angle);
            if (lbar_delta_angle > MAX_LIGHTBAR_DELTA) continue;

            // ------------防止出现带有中间灯条的误匹配----------------------------------
            double armor_area = (lbar->rect.size.area() + rbar->rect.size.area()) * 0.5;
            bool valid = true;
            auto temp_iter = rbar;
            std::advance(temp_iter, -1);
            if (lbar != temp_iter)
            {
                auto lbar_mid = lbar;
                int MaxY = 0, MinY = 0;
                MinY = int(std::min(lbar->rect.boundingRect().y, rbar->rect.boundingRect().y));
                MaxY = int(std::max(
                    lbar->rect.boundingRect().y + lbar->rect.boundingRect().height,
                    rbar->rect.boundingRect().y + rbar->rect.boundingRect().height));
                for (++lbar_mid; lbar_mid != rbar; lbar_mid++)
                {
                    if ((lbar_mid->rect.center.y < MaxY && lbar_mid->rect.center.y > MinY)
                        && (abs(lbar_mid->rect.size.height - armor_height) < armor_height / 5.0)
                        && abs(lbar_mid->rect.size.area() - armor_area) < armor_area / 5.0)
                    {
                        valid = false;
                        break;
                    }
                }
            }
            if (!valid) continue;

            Armor armor;
            armor.distance = armor_distance;
            armor.assignMemberVariables(*lbar, *rbar);
            armor.setNumberImg(img);
            armor.classifier_num = classifier(armor.number_img);
            armors.emplace_back(std::move(armor));
        }
    }
}

void Aimbot::gradeArmors(std::vector<Armor> &armors, int max_armor_num)
{
    if (armors.empty()) return;

    std::priority_queue<std::tuple<Armor, ErrorScore>, std::vector<std::tuple<Armor, ErrorScore>>, ErrorScoreComp> err_armors;

    ErrorScore error_score;

    for (auto &armor: armors)
    {
        // 距离误差分数
        error_score.dis_error_score = get_distance_score(armor.distance);
        error_score.sum_error_score += error_score.dis_error_score;
        // std::cout << "小孔成像距离：" << distance << std::endl;

        // 权重
        if (armor.isBigArmor())
        {
            error_score.ratio_error_score = 2 * fabs(armor.ratio - 4.09);
            error_score.sum_error_score += error_score.ratio_error_score;
        }
        else
        {
            error_score.ratio_error_score = 2 * fabs(armor.ratio - 2.363636);
            error_score.sum_error_score += error_score.ratio_error_score;
        }

        // 两灯条高度比造成的误差分数
        error_score.bar_len_ratio_error_score = (armor.bar_len_ratio - 1.0) * 10.0 / (MAX_LENGTH_RATIO - 1.0);
        error_score.sum_error_score += error_score.bar_len_ratio_error_score;

        // 两灯条倾斜度之差不过大，因为如果是正常的匹配的灯条角度差比较小，权重要稍微高一点
        error_score.bar_delta_angle_error_score = (armor.bar_delta_angle / (MAX_LIGHTBAR_DELTA)) * 15.0;
        error_score.sum_error_score += error_score.bar_delta_angle_error_score;

        // 装甲板倾斜角度误差分数
        error_score.armor_angle_error_score = fabs(armor.armor_angle) * 15.0 / MAX_ARMOR_ANGLE;
        error_score.sum_error_score += error_score.armor_angle_error_score;
        err_armors.push(std::make_tuple(armor, error_score));
    }
    armors.clear();
    while (!err_armors.empty() && armors.size() < max_armor_num)
    {
        armors.emplace_back(std::get<0>(err_armors.top()));
        err_armors.pop();
    }
}


double Aimbot::get_distance_score(double distance)
{
    double score = 0;
    if (distance < 2.0) score = 0;
    else if (distance >= 2.0 && distance < 3.0)
        score = distance;
    else if (distance >= 3.0 && distance < 4.0)
        score = 8.0 * (distance - 3) + 3;
    else if (distance >= 4.0 && distance < 5.0)
        score = 12.0 * (distance - 4.0) + 11.0;
    else
        score = 4 * distance;
    return score;
}

void
Aimbot::processImage(const cv::Mat &src, Armor::Color enemy_color, cv::Mat &processed_image)
{
    if (enemy_color == Armor::Color::BLUE)
    {
        thresh_img(src, processed_image, enemy_color, BLUE_BINARY_THRES, BLUE_CHANNEL_THRESH);
    }
    else
    {
        thresh_img(src, processed_image, enemy_color, BLUE_BINARY_THRES, RED_CHANNEL_THRESH);
    }
}

void Aimbot::adjustLightBarRect(cv::RotatedRect &rect)
{
    if (rect.size.width > rect.size.height)
    {
        rect = cv::RotatedRect(rect.center, cv::Size2f(rect.size.height, rect.size.width),
                               rect.angle + 90);
    }
}

void Aimbot::adjustLightBarAngle(float elli_angle, double &bar_angle)
{
    if (elli_angle > 90.0)
    { bar_angle = elli_angle - 180.0; }
    else
    {
        bar_angle = elli_angle;
    }
}

void Aimbot::judgeLightbarColor(Armor::Lightbar &lightbar, const cv::Mat &src)
{
    auto region = lightbar.rect.boundingRect();

    region.x -= std::max(2.0, region.width * 0.1);
    region.y -= std::max(2.0, region.height * 0.05);
    region.width += 2 * std::max(3.0, region.width * 0.1);
    region.height += 2 * std::max(3.0, region.height * 0.05);
    // roi
    region &= cv::Rect(0, 0, src.cols, src.rows);
    cv::Mat roi = src(region);

    int red_cnt = 0, blue_cnt = 0;

    for (int row = 0; row < (roi.rows / 2); row++)
    {
        for (int col = 0; col < roi.cols; col++)
        {
            red_cnt += roi.at<cv::Vec3b>(row, col)[2];
            blue_cnt += roi.at<cv::Vec3b>(row, col)[0];
        }
    }
    if (red_cnt > blue_cnt)
    { lightbar.color = Armor::Color::RED; }
    else
    {
        lightbar.color = Armor::Color::BLUE;
    }
}

void
Aimbot::thresh_img(const cv::Mat &src, cv::Mat &out, Armor::Color color, int gray_thresh, int color_channel_thresh)
{
    out.create(src.rows, src.cols, CV_8UC1);
    //每行循环处理
    for (int y = 0; y < src.rows; y++)
    {
        auto *cp = src.ptr<uchar>(y); //获取第y行像素的首地址
        auto *gp = out.ptr<uchar>(y);//获取第y行像素的首地址
        for (int x = 0; x < src.cols; x++)
        {
            *gp = (15 * cp[0] + 75 * cp[1] + 38 * cp[2]) >> 7;
            if (*gp > gray_thresh)
            {
                if (color == Armor::Color::BLUE) // Mat默认BGR
                {
                    if (cp[0] > color_channel_thresh)
                    {
                        *gp = 255;
                    }
                    else
                    {
                        *gp = 0;
                    }
                }
                else
                {
                    if (cp[2] > color_channel_thresh)
                    {
                        *gp = 255;
                    }
                    else
                    {
                        *gp = 0;
                    }
                }
            }
            else
            {
                *gp = 0;
            }
            cp += 3;
            gp++;
        }
    }
}

inline bool Aimbot::checkLightbarSize(const std::vector<Armor::Lightbar> &lightbars)
{
    return !(lightbars.empty() || lightbars.size() == 1);
}
